import base64
import traceback
import json
from collections import OrderedDict
from pyim.middleware import PacketReject
from pyim import config, logger
import hashlib
from pyim.packet import Packet

logger = logger.get_logger(__package__)


class WebSocketMiddleWare:
    WS_KEY = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

    def do_data_exchange(self, conn, data):
        conn.serializer = "websocket"
        logger.info("set the websocket as serializer")

    def encrypt_key(self, key):
        sha1 = hashlib.sha1()
        sha1.update((key + self.WS_KEY).encode("utf-8"))
        return base64.b64encode(sha1.digest()).decode("utf-8")

    def pre_process(self, connection, packet):
        try:
            if connection.serializer == 'http_parser':
                assert "method" in packet and packet['method'] == "GET"
                assert "uri" in packet
                assert "protocol" in packet
                assert "headers" in packet
                assert "Connection" in packet['headers'] and packet['headers']['Connection'] == 'Upgrade'
                assert "Upgrade" in packet['headers'] and packet['headers']['Upgrade'] == 'websocket'

                reply = Packet()
                reply["type"] = "response"
                reply["protocol"] = packet["protocol"]
                reply["version"] = packet["version"]
                reply["status_code"] = 101
                reply["status"] = "Switching Protocols"

                reply["headers"] = OrderedDict({
                    "Upgrade": "websocket",
                    "Connection": "Upgrade",
                    "Sec-WebSocket-Accept": self.encrypt_key(packet["headers"]['Sec-WebSocket-Key']),
                    "Sec-WebSocket-Location": "ws://%s:%d" % (config.bind_address[0],
                                                              config.bind_address[1])
                })

                logger.info("Websocket request is ok now")
                connection.route(reply)
                connection.on_sent = self.do_data_exchange
                raise PacketReject
            elif connection.serializer == 'websocket':
                if packet['opcode'] == 0x81:
                    import json
                    data = json.loads(packet['data'])
                    for k in data:
                        packet[k] = data[k]
                elif packet['opcode'] == 0x88:
                    connection.disconnect()

        except AssertionError:
            traceback.print_exc()
            return None
        except PacketReject:
            raise PacketReject
        except Exception as e:
            traceback.print_exc()
            logger.error(str(e))
            return None

    def post_process(self, connection, packet):
        logger.info("post process for websocket:" + connection.serializer)
        if connection.serializer == "websocket":
            data = {}
            for k in list(packet.keys()):
                data[k] = packet[k]
                del packet[k]
            packet['opcode'] = 129
            packet['data'] = json.dumps(data)
